import sounddevice
from scipy.io.wavfile import write
from pynput.keyboard import Listener, Key, KeyCode
import threading
import time
import numpy
import os
from threading import Event
from queue import Queue

from get_filenames import GetFilename

# 常用终端颜色代码
RESET = "\033[0m"      # 重置所有样式

# 常规颜色（较柔和）
BLACK = "\033[30m"     # 黑色
RED = "\033[31m"       # 红色
GREEN = "\033[32m"     # 绿色
YELLOW = "\033[33m"    # 黄色（比红色柔和，推荐）
BLUE = "\033[34m"      # 蓝色
MAGENTA = "\033[35m"   # 洋红色
CYAN = "\033[36m"      # 青蓝色（柔和，推荐）
WHITE = "\033[37m"     # 白色

# 亮色调（更鲜艳）
LIGHT_BLACK = "\033[90m"   # 亮黑（灰色）
LIGHT_RED = "\033[91m"     # 亮红
LIGHT_GREEN = "\033[92m"   # 亮绿
LIGHT_YELLOW = "\033[93m"  # 亮黄
LIGHT_BLUE = "\033[94m"    # 亮蓝
LIGHT_MAGENTA = "\033[95m" # 亮洋红
LIGHT_CYAN = "\033[96m"    # 亮青（柔和，推荐）
LIGHT_WHITE = "\033[97m"   # 亮白

class AudioRecorder:
    """音频录制类，支持自定义按键控制录音开始和停止"""
    
    def __init__(
            self, 
            is_recording_event: Event,

            stt_stopped_event: Event,
            nlp_stopped_event: Event,
            tts_stopped_event: Event,
            playmusic_stopped_event: Event,

            exit_event: Event,

            recording_filename_queue: Queue,
            playmusic_filename_queue: Queue,
            playeffect_filename_queue: Queue,

            samplerate=44100, 
            channels=1, 
            filename="", 
            record_key=Key.f8, 
            use_exit_key=True, 
            exit_key=Key.f8,
            exit_key2_pressed=Key.ctrl_l,
            enable_print=True,
            delete_recording=True,
            project_name="LLMSpeech"
            ):
        """
        record_key: 控制录音的按键，默认为F8
        自定义为F9键 record_key=Key.f9
        自定义为字母键，例如"r"键 record_key=KeyCode.from_char('r')
        """
        # event
        self.is_recording_event = is_recording_event
        self.stt_stopped_event = stt_stopped_event
        self.nlp_stopped_event = nlp_stopped_event
        self.tts_stopped_event = tts_stopped_event
        self.playmusic_stopped_event = playmusic_stopped_event
        self.exit_event = exit_event
        # 队列
        self.recording_filename_queue = recording_filename_queue
        self.playmusic_filename_queue = playmusic_filename_queue
        self.playeffect_filename_queue = playeffect_filename_queue
        # 音频参数
        self.samplerate = samplerate
        self.channels = channels
        self.output_file = filename
        self.use_get_filename = False
        if self.output_file:
            self.use_get_filename = False
        else:
            self.use_get_filename = True
        self.record_key = record_key  # 控制录音的按键
        self.exit_key = exit_key  # 退出程序的按键之一,默认也是 F8
        self.is_recording = False
        self.recording_data = None
        self.stream = None  # 音频流对象
        self.listener = None  # 键盘监听器
        self.recording_filename = GetFilename(r"temps/recording_", ".wav")
        self.use_exit_key = use_exit_key
        self.exit_key = exit_key
        self.enable_print = enable_print
        self.delete_recording = delete_recording
        self.filenames = []
        self.project_name = project_name
        # 组合按键跟踪
        self.exit_key2_pressed = exit_key2_pressed
        self.exit_pressed = False

        """时间控制"""
        # 上次录音结束时间
        self.last_recording_end_time = 0
        # 最短录音时长
        self.min_recording_duration = 1.0
        # 录音冷却时间
        self.cooldown_period = 1.0
        # 录音开始时间
        self.recording_start_time = 0

    def start_recording(self):
        """开始录音"""
        # 检查是否已冷却
        current_time = time.time()
        if current_time - self.last_recording_end_time < self.cooldown_period:
            if self.enable_print:
                print("录音冷却中，请稍后再试")
            return

        if not self.is_recording:
            if self.enable_print:
                # time.sleep(0.01)
                print()
                print(f"{RED}录音中... (松开{self._get_key_name().upper()}停止){RESET}")
            self.is_recording = True
            # 记录录音开始时间
            self.recording_start_time = time.time()
            # 通知其他线程录音开始
            self.is_recording_event.set()
            # 推送录音开始提示音
            self.playeffect_filename_queue.put(r"assets\waves\recording.wav")
            # 创建音频流
            self.stream = sounddevice.InputStream(
                samplerate=self.samplerate, 
                channels=self.channels, 
                dtype='int16'
            )
            self.recording_data = []
            self.stream.start()

            # 持续读取音频数据
            while self.is_recording:
                data, _ = self.stream.read(1024)  # 每次读取1024帧
                self.recording_data.append(data)
                time.sleep(0.01)  # 减少CPU占用

    def stop_recording(self):
        """停止录音并保存"""
        if self.is_recording:
            # 计算录音时长
            recording_duration = time.time() - self.recording_start_time
            if recording_duration < self.min_recording_duration:
                sleep_time = self.min_recording_duration - recording_duration
                if self.enable_print:
                    print(f"录音时长不足{self.min_recording_duration}秒，请注意")
                time.sleep(sleep_time)

            self.is_recording = False
            time.sleep(0.02)
            if self.stream:
                self.stream.stop()
                self.stream.close()

            # 等待其他线程处理完成
            self.stt_stopped_event.wait()
            self.nlp_stopped_event.wait()
            self.tts_stopped_event.wait()
            self.playmusic_stopped_event.wait()
            # 通知其他线程录音结束
            self.stt_stopped_event.clear()
            self.nlp_stopped_event.clear()
            self.tts_stopped_event.clear()
            self.playmusic_stopped_event.clear()

            # 通知其他线程停止录音
            self.is_recording_event.clear()

            # 推送录音结束提示音
            self.playeffect_filename_queue.put(r"assets\waves\recorded.wav")

            # 合并录音数据并保存
            if self.use_get_filename:
                self.output_file = self.recording_filename.get_filename()
            write(self.output_file, self.samplerate, numpy.concatenate(self.recording_data))

            self.recording_filename_queue.put(self.output_file)
            self.filenames.append(self.output_file)
            # 记录上次录音结束时间
            self.last_recording_end_time = time.time()

    def _get_key_name(self):
        """获取按键的名称，用于友好提示"""
        if isinstance(self.record_key, Key):
            return self.record_key.name
        elif isinstance(self.record_key, KeyCode):
            return self.record_key.char
        return str(self.record_key)

    def on_press(self, key):
        """按键按下事件处理"""
        try:
            # 更新ctrl状态
            if key == self.exit_key2_pressed:
                self.exit_pressed = True
                return
            
            # 判断组合键是否成立
            if self.exit_pressed and key == self.exit_key:
                print(f"\n正在退出 {LIGHT_GREEN}{self.project_name}{RESET}...")
                self.exit_event.set()
                if self.delete_recording:
                    for filename in self.filenames:
                        os.remove(filename)
                return False

            if key == self.record_key:
                threading.Thread(target=self.start_recording, daemon=True).start()
        except Exception as e:
            print(f"按键处理错误: {e}")

    def on_release(self, key):
        """按键松开事件处理"""
        try:
            if key == self.exit_key2_pressed:
                self.exit_pressed = False
                return
            if key == self.record_key:
                self.stop_recording()
        except Exception as e:
            print(f"释放处理错误: {e}")

    def _waiting_start(self):
        """等待各个模块启动成功"""
        self.nlp_stopped_event.wait()
        self.stt_stopped_event.wait()
        self.tts_stopped_event.wait()
        self.playmusic_stopped_event.wait()

        self.stt_stopped_event.clear()
        self.nlp_stopped_event.clear()
        self.tts_stopped_event.clear()
        self.playmusic_stopped_event.clear()

        self._start_success()
      
    def _start_success(self):
        # 定义颜色代码
        BOLD_LIGHT_GREEN = "\033[1;92m"
        # 消息内容
        message = f"🌟 {self.project_name} 启动成功！🌟"
        # 计算实际字符长度（中文算2个字符宽度）
        def str_width(s):
            return sum(2 if ord(c) > 127 else 1 for c in s)
        msg_width = str_width(message)
        total_width = msg_width + 4  # 两边各加2个空格
        # 顶部边框
        print(f"{LIGHT_BLUE}╔{'═' * total_width}╗{RESET}")
        # 内容行 - 使用居中对齐确保两侧空格相等
        print(f"{LIGHT_BLUE}║{RESET}  {BOLD_LIGHT_GREEN}{message}{RESET}  {LIGHT_BLUE}║{RESET}")
        # 底部边框
        print(f"{LIGHT_BLUE}╚{'═' * total_width}╝{RESET}")

    def start(self):
        """开始监听键盘事件"""
        self._waiting_start()

        key_name = self._get_key_name()
        if self.enable_print :
            print(f"按 {LIGHT_RED}{key_name.upper()}{RESET} 开始录音，松开 {LIGHT_RED}{key_name.upper()}{RESET} 停止录音", end="")
        if self.use_exit_key and self.enable_print:
            print(f", 按 {LIGHT_RED}{self.exit_key2_pressed.name.upper()} + {self.exit_key.name.upper()}{RESET} 退出程序")
            for _ in range(34):
                print("=", end="")
            print()
        else:
            print()
        
        with Listener(on_press=self.on_press, on_release=self.on_release) as self.listener:
            self.listener.join()


if __name__ == "__main__":
    pass